import openai

from utils.StaticObj import class_to_dict_list
class Message:
    role = "user"
    content = ""
    name = ''
    function_call = None

    def __init__(self, role, content, name=None, function_call=None):
        self.role = role
        self.content = content
        self.name = name
        self.function_call = function_call

    def json(self):
        dict = {
            'role': self.role,
            'content': self.content,
        }
        if self.name:
            dict['name'] = self.name
        if self.function_call:
            dict['function_call'] = self.function_call
        return dict


class GptQuery:
    chatModel = 'gpt-3.5-turbo-16k'
    contrastModel = 'ft:gpt-3.5-turbo-0613:sgtoone-ii::88g1dZzw'
    # contrastoneModel = 'ft:gpt-3.5-turbo-0613:sgtoone-ii::8ADijPSf'
    contrastoneModel = 'ft:gpt-3.5-turbo-0613:sgtoone-ii::8D1bJowT'
    function_call = 'auto'
    temperature = 0.0
    stream = True
    constream = False

    systemMessage = {}
    messages = []

    def __init__(self, systemMessage=None, messages=None):
        self.systemMessage = systemMessage
        self.messages = []
        if systemMessage:
            self.messages.append(self.systemMessage.json())
        if messages:
            self.messages.extend(class_to_dict_list(self.shaer(messages)))

    def shaer(self,messages):
        if messages.__len__() >17:
            return messages[17:]
        else:
            return messages

    def query(self, message, functions=None):
        self.messages.append(message.json())
        # print("messages length:",self.messages.__len__())
        if functions:
            return openai.ChatCompletion.create(
                model=self.chatModel,
                messages=self.messages,
                temperature=self.temperature,
                functions=functions,
                stream=self.stream,
                function_call=self.function_call,  # auto is default, but we'll be explicit
            )
        else:
            return openai.ChatCompletion.create(
                model=self.chatModel,
                messages=self.messages,
                temperature=self.temperature,
                stream=self.stream
            )

    def contrast(self, message, functions=None):
        self.messages.append(message)
        # print("messages length:",self.messages.__len__())
        if functions:
            return openai.ChatCompletion.create(
                model=self.contrastoneModel,
                messages=self.messages,
                temperature=self.temperature,
                functions=functions,
                stream=self.stream,
                function_call=self.function_call,  # auto is default, but we'll be explicit
            )
        else:
            return openai.ChatCompletion.create(
                model=self.contrastoneModel,
                messages=self.messages
            )


